using System.Collections.Concurrent;
using DocumentCreationSystem.Models;
using Microsoft.Extensions.Logging;

namespace DocumentCreationSystem.Services
{
    /// <summary>
    /// 内置的轻量记忆系统：内存向量+倒排，演示用途，可切换到 Ai00MemAdapter
    /// </summary>
    public class BuiltInMemoryService : IMemoryService
    {
        private readonly ILogger<BuiltInMemoryService> _logger;
        private readonly IVectorService _vectorService;
        private readonly ConcurrentDictionary<string, MemoryRecord> _store = new();
        private float _threshold = 0.7f;

        public BuiltInMemoryService(ILogger<BuiltInMemoryService> logger, IVectorService vectorService)
        { _logger = logger; _vectorService = vectorService; }

        public async Task<bool> InitializeAsync(MemoryConfig config)
        {
            _threshold = config.SimilarityThreshold;
            await _vectorService.InitializeAsync();
            return true;
        }

        public async Task<string> AddMemoryAsync(string content, string? userId = null, string? sessionId = null, Dictionary<string, object>? attributes = null)
        {
            var id = Guid.NewGuid().ToString("N");
            var emb = await _vectorService.EmbedTextAsync(content);
            _store[id] = new MemoryRecord
            {
                Id = id,
                Content = content,
                Attributes = attributes ?? new(),
                Embedding = emb
            };
            return id;
        }

        public async Task<List<MemoryResult>> QueryAsync(string query, string? userId = null, int topK = 5, float threshold = 0.7f)
        {
            var qEmb = await _vectorService.EmbedTextAsync(query);
            var th = threshold > 0 ? threshold : _threshold;
            var vecs = _store.Values
                .Select(v => new { v.Id, v.Content, v.Attributes, Score = CosSim(qEmb, v.Embedding) })
                .Where(x => x.Score >= th)
                .OrderByDescending(x => x.Score)
                .Take(topK > 0 ? topK : 5)
                .Select(x => new MemoryResult { Id = x.Id, Score = x.Score, Content = x.Content, Attributes = x.Attributes })
                .ToList();
            return vecs;
        }

        private static float CosSim(float[] a, float[] b)
        {
            double dot = 0, na = 0, nb = 0;
            for (int i = 0; i < Math.Min(a.Length, b.Length); i++) { dot += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i]; }
            if (na == 0 || nb == 0) return 0;
            return (float)(dot / (Math.Sqrt(na) * Math.Sqrt(nb)));
        }

        private class MemoryRecord
        {
            public string Id { get; set; } = string.Empty;
            public string Content { get; set; } = string.Empty;
            public Dictionary<string, object> Attributes { get; set; } = new();
            public float[] Embedding { get; set; } = Array.Empty<float>();
        }
    }
}

